import os
import pandas as pd
from pymongo import MongoClient


def loc_collection(db_name, collection_name):
    conn = MongoClient("mongodb://192.168.3.179:27017")
    coll_ = conn[db_name][collection_name]
    return coll_


def loc_zq_collection(db_name, collection_name):
    conn = MongoClient("mongodb://192.168.1.9:27017")
    coll_ = conn[db_name][collection_name]
    return coll_


def loc_zq_db(db_name):
    return MongoClient("mongodb://192.168.1.9:27017")[db_name]


def get_data():
    db = loc_zq_collection()


def get_df_by_dict(s: dict):
    re = pd.DataFrame(s['data'], index=s['index'], columns=s['columns'])
    re.sort_index(ascending=True, inplace=True)
    # print(re.columns)
    return re.iloc[:, 0]


def get_data_by_db_name(db_name):
    meta_data = list(loc_zq_collection(db_name, 'metadata').find())
    map_id = {s['__id']: s['name'] for s in meta_data}
    data_li = list(loc_zq_collection(db_name, 'data').find())
    assert len(data_li) > 0
    df_data = pd.concat([get_df_by_dict(s) for s in data_li], axis=1)
    df_data.columns = [map_id[s] for s in df_data.columns]
    return df_data


def ini_all_x_data():
    conn = MongoClient("mongodb://192.168.1.9:27017")
    db_li = conn.list_database_names()
    filter_db_li = [s for s in db_li if 'x_' in s]
    if not os.path.exists('.input'):
        os.mkdir('input')
    for db_name in filter_db_li:
        try:
            data = get_data_by_db_name(db_name)
            data.to_pickle(f"./input/{db_name}.pkl")
            print(f'download {db_name} sucesss !!!')
        except:
            print(f'download {db_name} failed !!!')


if __name__ == "__main__":
    ini_all_x_data()
